from prefect import task
import pandas as pd
from prefect import get_run_logger

from application.db.mysql_db.info import ResourceSourceDict
from application.tasks.base_task import BaseTask


class DataTypeTask(BaseTask):
    tag_code = "info_nsfc"  # 标签类型

    def run(self, df: pd.DataFrame):
        """
            清洗并计算文章信息的统计结果

            :param df: pandas.DataFrame
                输入的文献信息数据表，必须包含 `article_info` 字段
            :return: pandas.DataFrame
                返回清洗并增加检查字段后的 DataFrame
            """

        resorce_list = ResourceSourceDict.select(ResourceSourceDict.source_main_link, ResourceSourceDict.source_id)
        resorce_dict = {
            resorce.source_main_link: resorce.source_id
            for resorce in resorce_list
        }  # 资源字典
        # logger = get_run_logger()
        # logger.info("Computing statistics")
        df['tag_code'] = self.tag_code  # 标签类型
        df['source_id'] = df['metadata'].apply(lambda x: resorce_dict.get(x['allowed_domains'][0], None))  # 资源id
        df['metadata'] = df.apply(lambda x: {"info_author": x['info_author']}, axis=1)  # 元数据
        df['information_name'] = df['info_name'].apply(lambda x: {"zh": x})  # 信息名称
        return df
